import random
import json
import logging
import traceback
import hashlib
import hmac
from django.http import JsonResponse
from django.views.decorators.csrf import csrf_exempt
from qiskit_aer import Aer
from qiskit import QuantumCircuit
from .models import qkd_keys, temp_storage
from django.conf import settings

# 从settings获取QKD配置
KEY_LENGTH = getattr(settings, 'QKD_KEY_LENGTH', 128)
BATCH_SIZE = getattr(settings, 'QKD_BATCH_SIZE', 10)  # 减小批处理大小为10

# 配置日志
logger = logging.getLogger(__name__)


def adapt_key_for_aes128(qkd_key):
    #将QKD密钥调整为AES-128兼容的长度 (128位/16字节/32个十六进制字符)
    key_len_bits = len(qkd_key) * 4  # 十六进制字符串长度 * 4

    if key_len_bits < 128:
        # 如果密钥太短，使用HKDF扩展到128位
        salt = b'AES-QKD-SALT'
        key_material = bytes.fromhex(qkd_key)
        prk = hmac.new(salt, key_material, hashlib.sha256).digest()
        aes_key = prk[:16]  # 取前16字节(128位)
        return aes_key.hex()
    else:
        # 如果密钥足够长或更长，截取到128位
        return qkd_key[:32]  # 取前32个十六进制字符(128位)


def alice_prepare_qubits(n):
    """生成Alice的量子比特和基"""
    alice_bits = [random.randint(0, 1) for _ in range(n)]
    alice_basis = [random.choice(['+', 'x']) for _ in range(n)]

    # 模拟Alice在量子信道上准备量子比特
    try:
        simulator = Aer.get_backend('aer_simulator')
    except:
        # 如果aer_simulator不可用，尝试使用基本模拟器
        try:
            simulator = Aer.get_backend('statevector_simulator')
        except:
            # 如果所有模拟器都不可用，记录错误
            logger.error("无法加载量子模拟器，请检查Qiskit安装")
            raise Exception("量子模拟器加载失败")

    measurements = []

    # 分批处理量子比特
    for batch_start in range(0, n, BATCH_SIZE):
        batch_end = min(batch_start + BATCH_SIZE, n)
        batch_size = batch_end - batch_start

        # 创建量子电路
        qc = QuantumCircuit(batch_size, batch_size)

        # 准备量子比特
        for i in range(batch_size):
            idx = batch_start + i
            if alice_bits[idx] == 1:
                qc.x(i)  # 设置为|1⟩状态

            if alice_basis[idx] == 'x':
                qc.h(i)  # 如果使用X基，应用Hadamard门

        # 添加测量门
        qc.measure(range(batch_size), range(batch_size))

        # 模拟测量结果
        try:
            job = simulator.run(qc, shots=1, memory=True)
            result = job.result()
            memory_result = result.get_memory()[0]
            batch_measurements = list(map(int, memory_result))
        except Exception as e:
            # 如果量子模拟失败，直接抛出异常
            logger.error(f"量子模拟失败: {str(e)}")
            raise Exception(f"量子模拟失败: {str(e)}")

        measurements.extend(batch_measurements)

    return {
        "bits": alice_bits,
        "basis": alice_basis,
        "measurements": measurements
    }


def calculate_shared_key(alice_data, bob_basis):
    #计算Alice的共享密钥
    alice_basis = alice_data["basis"]
    alice_bits = alice_data["bits"]

    # 保留使用相同基的比特
    shared_key = [
        alice_bits[i] for i in range(len(bob_basis))
        if alice_basis[i] == bob_basis[i]
    ]

    # 转换为十六进制
    hex_key = ''.join(format(int(''.join(map(str, shared_key[i:i + 8])), 2), '02x')
                      for i in range(0, len(shared_key), 8) if i + 8 <= len(shared_key))

    return hex_key


@csrf_exempt
def start_qkd(request):
    #开始QKD过程 - 服务器作为Alice
    if request.method != 'POST':
        return JsonResponse({"status": "error", "message": "只允许POST请求"}, status=405)

    try:
        data = json.loads(request.body)
        username = data.get("username")

        if not username:
            return JsonResponse({"status": "error", "message": "用户名不能为空"}, status=400)

        # 记录详细日志
        logger.info(f"开始为用户 {username} 生成QKD密钥")

        # 生成并准备量子比特
        qubits_count = KEY_LENGTH * 2  # 生成2倍长度，因为部分基不匹配会被丢弃

        try:
            alice_data = alice_prepare_qubits(qubits_count)
            logger.info(f"成功生成量子比特: {len(alice_data['bits'])} 位")
        except Exception as qe:
            logger.error(f"量子比特生成错误: {str(qe)}")
            return JsonResponse({"status": "error", "message": f"量子比特生成错误: {str(qe)}"}, status=500)

        # 使用临时存储而非会话存储，避免序列化问题
        storage_key = f'alice_data_{username}'
        temp_storage[storage_key] = alice_data

        # 返回Alice的数据给Bob
        return JsonResponse({
            "status": "success",
            "message": "QKD会话已开始",
            "alice_data": {
                "basis": alice_data["basis"],
                "bits": alice_data["bits"]  # 发送原始比特而不是测量结果
            }
        })

    except Exception as e:
        logger.error(f"QKD启动出错: {str(e)}\n{traceback.format_exc()}")
        return JsonResponse({"status": "error", "message": f"QKD启动出错: {str(e)}"}, status=500)


@csrf_exempt
def complete_qkd(request):
   #完成QKD过程 接收Bob的基并计算共享密钥
    if request.method != 'POST':
        return JsonResponse({"status": "error", "message": "只允许POST请求"}, status=405)

    try:
        data = json.loads(request.body)
        username = data.get("username")
        bob_basis = data.get("bob_basis")

        if not username or not bob_basis:
            return JsonResponse({"status": "error", "message": "缺少必要参数"}, status=400)

        # 从临时存储获取Alice数据
        storage_key = f'alice_data_{username}'

        if storage_key not in temp_storage:
            return JsonResponse({"status": "error", "message": "QKD会话已过期或不存在"}, status=400)

        alice_data = temp_storage[storage_key]

        # 计算共享密钥
        raw_key = calculate_shared_key(alice_data, bob_basis)
        logger.info(f"用户 {username} 的原始QKD密钥生成成功，长度: {len(raw_key) * 4} 位")
        logger.info(f"用户 {username} 的原始密钥: {raw_key}")

        # 将密钥调整为AES-128兼容长度
        aes_key = adapt_key_for_aes128(raw_key)
        logger.info(f"用户 {username} 的AES-128适配密钥，长度: 128 位")
        logger.info(f"用户 {username} 的AES适配密钥: {aes_key}")

        # 存储适配后的密钥
        qkd_keys[username] = aes_key

        # 清除临时数据以释放内存
        del temp_storage[storage_key]

        return JsonResponse({
            "status": "success",
            "message": "QKD过程已完成",
            "key": aes_key,
            "raw_key": raw_key,
            "raw_key_length": len(raw_key) * 4
        })

    except Exception as e:
        logger.error(f"QKD完成过程出错: {str(e)}\n{traceback.format_exc()}")
        return JsonResponse({"status": "error", "message": f"QKD完成过程出错: {str(e)}"}, status=500)


@csrf_exempt
def verify_key(request):
    #验证用户是否已有QKD密钥
    if request.method != 'POST':
        return JsonResponse({"status": "error", "message": "只允许POST请求"}, status=405)

    try:
        data = json.loads(request.body)
        username = data.get("username")

        if not username:
            return JsonResponse({"status": "error", "message": "用户名不能为空"}, status=400)

        has_key = username in qkd_keys

        return JsonResponse({
            "status": "success",
            "has_key": has_key,
            "message": "用户已有QKD密钥" if has_key else "用户没有QKD密钥"
        })

    except Exception as e:
        logger.error(f"验证QKD密钥出错: {str(e)}")
        return JsonResponse({"status": "error", "message": f"验证QKD密钥出错: {str(e)}"}, status=500)